import glob
import os
import cv2
import numpy as np
import random

# 设置数据根目录
rgbt_cc_root = './data/DroneRGBT/'

# 定义数据集类型
datasets = ['train', 'val', 'test']

# 设置随机种子以确保可重复性
np.random.seed(0)
random.seed(0)

# 目标尺寸
TARGET_WIDTH = 672
TARGET_HEIGHT = 448
CROP_SIZE = 224  # 如果需要裁剪的话

def process_dataset(dataset_type):
    """处理指定类型的数据集"""
    # 创建输出目录
    input_path = os.path.join(rgbt_cc_root, dataset_type)
    output_path = os.path.join(rgbt_cc_root, f'new_{dataset_type}_224')
    
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    
    # 获取所有RGB图像路径
    img_paths = glob.glob(os.path.join(input_path, '*RGB.jpg'))
    img_paths.sort()
    
    print(f"Processing {dataset_type} set with {len(img_paths)} images...")
    
    for img_path in img_paths:
        # 读取RGB图像、热红外图像和标注
        img_data = cv2.imread(img_path)
        t_data = cv2.imread(img_path.replace('_RGB', '_T'))
        gt_data = np.load(img_path.replace('_RGB.jpg', '_GT.npy'))
        
        # 计算缩放比例
        if img_data.shape[1] >= img_data.shape[0]:  # 宽度大于或等于高度
            rate_w = TARGET_WIDTH / img_data.shape[1]
            rate_h = TARGET_HEIGHT / img_data.shape[0]
        else:  # 高度大于宽度
            rate_w = TARGET_WIDTH / img_data.shape[0]
            rate_h = TARGET_HEIGHT / img_data.shape[1]
        
        # 调整图像大小
        img_resized = cv2.resize(img_data, (0, 0), fx=rate_w, fy=rate_h)
        t_resized = cv2.resize(t_data, (0, 0), fx=rate_w, fy=rate_h)
        
        # 调整标注坐标
        gt_data[:, 0] = gt_data[:, 0] * rate_w
        gt_data[:, 1] = gt_data[:, 1] * rate_h
        
        # 生成输出路径
        base_name = os.path.basename(img_path)
        output_img_path = os.path.join(output_path, base_name)
        output_t_path = os.path.join(output_path, base_name.replace('_RGB', '_T'))
        output_gt_path = os.path.join(output_path, base_name.replace('_RGB.jpg', '_GT.npy'))
        
        # 保存处理后的数据
        cv2.imwrite(output_img_path, img_resized)
        cv2.imwrite(output_t_path, t_resized)
        np.save(output_gt_path, gt_data)
        
        print(f"Processed: {base_name}")

# 处理所有数据集
for dataset in datasets:
    process_dataset(dataset)

print("All datasets processed successfully!")